package org.chen.aao.gateway.filter;

import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.core.Ordered;
import org.springframework.data.redis.core.ReactiveRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.http.HttpStatus;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

import java.util.Arrays;
import java.util.Collections;
import java.util.List;

@Slf4j
@Component
@RequiredArgsConstructor
public class RateLimitFilter implements GlobalFilter, Ordered {
    
    private final ReactiveRedisTemplate<String, String> redisTemplate;
    
    // Lua脚本用于原子性地检查和更新限流计数
    private static final String LUA_SCRIPT = 
            "local key = KEYS[1]\n" +
            "local limit = tonumber(ARGV[1])\n" +
            "local window = tonumber(ARGV[2])\n" +
            "local current = redis.call('incr', key)\n" +
            "if current == 1 then\n" +
            "    redis.call('expire', key, window)\n" +
            "end\n" +
            "if current > limit then\n" +
            "    return 0\n" +
            "else\n" +
            "    return 1\n" +
            "end";

    private static final RedisScript<Long> RATE_LIMIT_SCRIPT = new DefaultRedisScript<>(LUA_SCRIPT, Long.class);
    
    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        ServerHttpRequest request = exchange.getRequest();
        
        // 获取用户ID或IP作为限流key
        String userId = request.getHeaders().getFirst("X-User-Id");
        String rateLimitKey;
        
        if (userId != null && !userId.isEmpty()) {
            rateLimitKey = "rate_limit:user:" + userId;
        } else {
            // 未登录用户使用IP限流
            String ip = getClientIp(request);
            rateLimitKey = "rate_limit:ip:" + ip;
        }
        
        // 设置限流参数（每分钟60次请求）
        int limit = userId != null ? 100 : 30; // 登录用户100次/分钟，未登录30次/分钟
        int window = 60; // 时间窗口（秒）

        List<String> args = Arrays.asList(String.valueOf(limit), String.valueOf(window));

        // 正确地响应式链
        return redisTemplate.execute(
                        RATE_LIMIT_SCRIPT,
                        Collections.singletonList(rateLimitKey),
                        args
                ).next() // 返回 Mono<Long>，但可能为空
                .defaultIfEmpty(0L) // 如果 Mono 为空，则提供默认值 0L
                .flatMap(result -> {
                    // 现在 result 保证有值 (来自 Redis 或默认值)
                    if (result == 1L) {
                        // 请求通过
                        return chain.filter(exchange);
                    } else {
                        // 请求被限流 (无论是 Redis 返回 0L 还是默认值 0L)
                        return handleRateLimitExceeded(exchange.getResponse());
                    }
                }).onErrorResume(error -> {
                    // 如果在执行过程中发生任何其他错误（如序列化问题），记录日志并放行
                    log.error("Rate limit check error: {}", error.getMessage(), error);
                    return chain.filter(exchange);
                });
    }
    
    private String getClientIp(ServerHttpRequest request) {
        List<String> xForwardedFor = request.getHeaders().get("X-Forwarded-For");
        if (xForwardedFor != null && !xForwardedFor.isEmpty()) {
            return xForwardedFor.get(0).split(",")[0].trim();
        }
        
        List<String> xRealIp = request.getHeaders().get("X-Real-IP");
        if (xRealIp != null && !xRealIp.isEmpty()) {
            return xRealIp.get(0);
        }
        
        return request.getRemoteAddress() != null ? 
               request.getRemoteAddress().getAddress().getHostAddress() : "unknown";
    }
    
    private Mono<Void> handleRateLimitExceeded(ServerHttpResponse response) {
        response.setStatusCode(HttpStatus.TOO_MANY_REQUESTS);
        response.getHeaders().add("X-RateLimit-Retry-After", "60");
        return response.setComplete();
    }
    
    @Override
    public int getOrder() {
        return -99;
    }
}